"""
Differential operators for spherically symmetric grids.

This module implements differential operators for spherical coordinate systems
with full spherical symmetry (radial dependence only). The key challenge is
proper handling of the r=0 singularity.

Mathematical form of operators:
- Laplacian: ∇²u = d²u/dr² + (2/r)du/dr
- Gradient: ∇u = du/dr
- Vector Laplacian: ∇²v = ∇²vᵣ - (2/r²)vᵣ (radial component)

Features:
- Proper treatment of r=0 singularity using L'Hôpital's rule
- Conservative and non-conservative formulations
- Support for Dirichlet, Neumann, and Robin boundary conditions
- High accuracy near the origin
- Mass-conserving discretizations
"""

using LinearAlgebra
using SparseArrays

include("common.jl")

# Export all spherical operator functions
export make_laplace_spherical, make_gradient_spherical, make_divergence_spherical
export make_vector_laplace_spherical, make_poisson_solver_spherical
export make_conservative_laplace_spherical

"""
    make_laplace_spherical(grid::SphericalSymGrid; conservative=false) -> OperatorType

Create Laplacian operator for spherically symmetric grid.

The spherical Laplacian in radial coordinates is:
∇²u = d²u/dr² + (2/r)du/dr

At r=0, using L'Hôpital's rule: (2/r)du/dr → 2d²u/dr²
So the effective Laplacian at r=0 is: ∇²u = 3d²u/dr²

Args:
    grid: Spherically symmetric grid
    conservative: Whether to use conservative (mass-preserving) formulation

Returns:
    Function that applies spherical Laplacian
"""
function make_laplace_spherical(grid::SphericalSymGrid; conservative=false)
    if conservative
        return make_conservative_laplace_spherical(grid)
    else
        return make_standard_laplace_spherical(grid)
    end
end

"""
    make_standard_laplace_spherical(grid::SphericalSymGrid) -> OperatorType

Create standard spherical Laplacian operator.
"""
function make_standard_laplace_spherical(grid::SphericalSymGrid)
    n = grid.num_r
    dr = spacing(grid)[1]
    rs = coordinate_arrays(grid)[1]  # radial coordinates
    r_min = extent(grid)[1][1]
    
    dr2_inv = 1.0 / (dr * dr)
    dr_inv = 1.0 / (2.0 * dr)
    
    function laplace_spherical(field_data, out)
        # Handle r=0 point specially
        if abs(r_min) < 1e-12  # r_min ≈ 0
            # At r=0: ∇²u = 6(u[2] - u[1])/dr² using L'Hôpital's rule
            out[1] = 6.0 * (field_data[2] - field_data[1]) * dr2_inv
        else
            # Regular interior point formula even for first point
            i = 1
            r = rs[i]
            d2u_dr2 = (field_data[i+1] - 2*field_data[i] + field_data[i-1]) * dr2_inv
            du_dr = (field_data[i+1] - field_data[i-1]) * dr_inv
            out[i] = d2u_dr2 + (2.0/r) * du_dr
        end
        
        # Interior points (standard formula)
        for i in 2:(n-1)
            r = rs[i]
            d2u_dr2 = (field_data[i+1] - 2*field_data[i] + field_data[i-1]) * dr2_inv
            du_dr = (field_data[i+1] - field_data[i-1]) * dr_inv
            out[i] = d2u_dr2 + (2.0/r) * du_dr
        end
        
        # Boundary point (will be handled by boundary conditions)
        out[n] = 0.0
        
        return out
    end
end

"""
    make_conservative_laplace_spherical(grid::SphericalSymGrid) -> OperatorType

Create conservative spherical Laplacian operator.

The conservative form discretizes: ∇²u = (1/r²)d/dr(r²du/dr)
This ensures exact mass conservation for diffusion problems.
"""
function make_conservative_laplace_spherical(grid::SphericalSymGrid)
    n = grid.num_r
    dr = spacing(grid)[1]
    rs = coordinate_arrays(grid)[1]
    r_min, r_max = extent(grid)[1]
    
    # Create shell boundaries for conservative discretization
    rl = rs .- dr/2  # inner shell radii
    rh = rs .+ dr/2  # outer shell radii
    
    # Ensure proper boundaries
    rl[1] = r_min
    rh[end] = r_max
    
    # Calculate shell volumes: V = (4π/3)(r_h³ - r_l³)
    # For conservative operator, we need V = (r_h³ - r_l³)/3
    volumes = (rh.^3 - rl.^3) / 3
    
    # Flux weighting factors
    factor_l = rl.^2 ./ (dr * volumes)
    factor_h = rh.^2 ./ (dr * volumes)
    
    function conservative_laplace_spherical(field_data, out)
        for i in 1:n
            if i == 1
                # Special handling at r=0
                # Flux from center is zero due to symmetry
                flux_h = factor_h[i] * (field_data[i+1] - field_data[i])
                out[i] = flux_h  # No flux_l term at center
            elseif i == n
                # Boundary point
                flux_l = factor_l[i] * (field_data[i] - field_data[i-1])
                out[i] = -flux_l  # Will be modified by boundary conditions
            else
                # Interior points
                flux_h = factor_h[i] * (field_data[i+1] - field_data[i])
                flux_l = factor_l[i] * (field_data[i] - field_data[i-1])
                out[i] = flux_h - flux_l
            end
        end
        
        return out
    end
end

"""
    make_gradient_spherical(grid::SphericalSymGrid; method="central") -> OperatorType

Create gradient operator for spherically symmetric grid.

In spherical symmetry, gradient is simply du/dr in radial direction.

Args:
    grid: Spherically symmetric grid  
    method: Finite difference method ("central", "forward", "backward")

Returns:
    Function that computes radial gradient
"""
function make_gradient_spherical(grid::SphericalSymGrid; method="central")
    n = grid.num_r
    dr = spacing(grid)[1]
    
    if method == "central"
        dr_inv = 1.0 / (2.0 * dr)
        
        function gradient_central(field_data, out)
            # Forward difference at r=0
            out[1] = (field_data[2] - field_data[1]) / dr
            
            # Central differences for interior points
            for i in 2:(n-1)
                out[i] = (field_data[i+1] - field_data[i-1]) * dr_inv
            end
            
            # Backward difference at boundary
            out[n] = (field_data[n] - field_data[n-1]) / dr
            
            return out
        end
        
    elseif method == "forward"
        dr_inv = 1.0 / dr
        
        function gradient_forward(field_data, out)
            for i in 1:(n-1)
                out[i] = (field_data[i+1] - field_data[i]) * dr_inv
            end
            out[n] = 0.0  # Boundary
            return out
        end
        
    elseif method == "backward"
        dr_inv = 1.0 / dr
        
        function gradient_backward(field_data, out)
            out[1] = 0.0  # Boundary
            for i in 2:n
                out[i] = (field_data[i] - field_data[i-1]) * dr_inv
            end
            return out
        end
        
    else
        error("Unsupported gradient method: $method")
    end
end

"""
    make_divergence_spherical(grid::SphericalSymGrid) -> OperatorType

Create divergence operator for vector fields in spherical coordinates.

For radially symmetric vectors: ∇·v = (1/r²)d/dr(r²vᵣ)

This is the same as the conservative form of scalar derivative.
"""
function make_divergence_spherical(grid::SphericalSymGrid)
    n = grid.num_r
    dr = spacing(grid)[1]
    rs = coordinate_arrays(grid)[1]
    
    function divergence_spherical(vector_field, out)
        # vector_field should have radial component only: vector_field[:, 1]
        vr = view(vector_field, :, 1)
        
        # Interior points: (1/r²)d/dr(r²vᵣ)
        for i in 2:(n-1)
            r = rs[i]
            # d/dr(r²vᵣ) using central differences
            flux_h = (rs[i] + dr/2)^2 * (vr[i+1] + vr[i])/2
            flux_l = (rs[i] - dr/2)^2 * (vr[i] + vr[i-1])/2
            out[i] = (flux_h - flux_l) / (dr * r^2)
        end
        
        # Special handling at r=0
        out[1] = 3.0 * (vr[2] - vr[1]) / dr  # Using L'Hôpital's rule
        
        # Boundary
        out[n] = 0.0
        
        return out
    end
end

"""
    make_vector_laplace_spherical(grid::SphericalSymGrid) -> OperatorType

Create vector Laplacian for spherically symmetric systems.

For a vector field in spherical coordinates with only radial component:
∇²v = (∇²vᵣ - 2vᵣ/r²)êᵣ

The angular components vanish due to symmetry.
"""
function make_vector_laplace_spherical(grid::SphericalSymGrid)
    scalar_laplace = make_laplace_spherical(grid)
    n = grid.num_r
    rs = coordinate_arrays(grid)[1]
    
    function vector_laplace_spherical(vector_field, out)
        # Only radial component is non-zero
        vr = view(vector_field, :, 1)
        
        # Apply scalar Laplacian to radial component
        temp = similar(vr)
        scalar_laplace(vr, temp)
        
        # Subtract correction term 2vᵣ/r²
        for i in 1:n
            r = rs[i]
            if r > 1e-12  # Avoid division by zero
                out[i, 1] = temp[i] - 2.0 * vr[i] / (r * r)
            else
                # At r=0, the correction term vanishes
                out[i, 1] = temp[i]
            end
        end
        
        return out
    end
end

"""
    make_poisson_solver_spherical(grid::SphericalSymGrid, bcs) -> Function

Create Poisson equation solver for spherical grid: ∇²u = f

Uses sparse matrix formulation with proper boundary condition handling.
"""
function make_poisson_solver_spherical(grid::SphericalSymGrid, bcs)
    return make_general_poisson_solver(grid, bcs)
end

"""
    make_laplace_matrix_spherical(grid::SphericalSymGrid, bcs) -> Tuple{SparseMatrixCSC, Vector}

Create sparse Laplacian matrix for spherical grid with boundary conditions.
"""
function make_laplace_matrix_spherical(grid::SphericalSymGrid, bcs)
    n = grid.num_r
    dr = spacing(grid)[1]
    rs = coordinate_arrays(grid)[1]
    r_min = extent(grid)[1][1]
    
    # Initialize sparse matrix components
    I_vals = Int[]
    J_vals = Int[]  
    V_vals = Float64[]
    rhs_vector = zeros(n)
    
    dr2_inv = 1.0 / (dr * dr)
    dr_inv = 1.0 / (2.0 * dr)
    
    for i in 1:n
        if i == 1 && abs(r_min) < 1e-12  # r=0 case
            # Special stencil at r=0: 6/dr²(-u₁ + u₂)
            push!(I_vals, i)
            push!(J_vals, i)  
            push!(V_vals, -6.0 * dr2_inv)
            
            push!(I_vals, i)
            push!(J_vals, i+1)
            push!(V_vals, 6.0 * dr2_inv)
            
        elseif i == n  # Boundary point
            # Handle boundary condition
            bc = bcs[1]  # Assuming single BC for radial direction
            if isa(bc, DirichletBC)
                # u_n = bc.value
                push!(I_vals, i)
                push!(J_vals, i)
                push!(V_vals, 1.0)
                rhs_vector[i] = bc.value
            elseif isa(bc, NeumannBC)
                # du/dr|_boundary = bc.value
                push!(I_vals, i)
                push!(J_vals, i)
                push!(V_vals, 1.0/dr)
                push!(I_vals, i)  
                push!(J_vals, i-1)
                push!(V_vals, -1.0/dr)
                rhs_vector[i] = bc.value
            end
            
        else  # Interior points
            r = rs[i]
            
            # Central point: -2/dr² term
            push!(I_vals, i)
            push!(J_vals, i)
            push!(V_vals, -2.0 * dr2_inv)
            
            # Forward difference: (1/dr² + 1/(r*dr)) term  
            push!(I_vals, i)
            push!(J_vals, i+1)
            push!(V_vals, dr2_inv + dr_inv/r)
            
            # Backward difference: (1/dr² - 1/(r*dr)) term
            push!(I_vals, i)
            push!(J_vals, i-1)  
            push!(V_vals, dr2_inv - dr_inv/r)
        end
    end
    
    matrix = sparse(I_vals, J_vals, V_vals, n, n)
    return matrix, rhs_vector
end

# Make the matrix function available to the general solver
make_laplace_matrix(grid::SphericalSymGrid, bcs) = make_laplace_matrix_spherical(grid, bcs)

"""
    spherical_volume_element(r, dr) -> Float64

Calculate volume element for spherical shell: 4πr²dr
"""
function spherical_volume_element(r, dr)
    return 4π * r * r * dr
end

"""
    spherical_surface_element(r) -> Float64

Calculate surface element for spherical surface: 4πr²
"""
function spherical_surface_element(r)
    return 4π * r * r
end

"""
    apply_spherical_bc!(operator_matrix, rhs_vector, bcs, grid)

Apply boundary conditions specifically for spherical grids.
"""
function apply_spherical_bc!(operator_matrix, rhs_vector, bcs, grid)
    n = grid.num_r
    bc = bcs[1]  # Single radial boundary condition
    
    if isa(bc, DirichletBC)
        # Set last row to identity
        operator_matrix[n, :] .= 0.0
        operator_matrix[n, n] = 1.0
        rhs_vector[n] = bc.value
        
    elseif isa(bc, NeumannBC)
        # Modify last row for gradient boundary condition
        dr = spacing(grid)[1]
        operator_matrix[n, :] .= 0.0
        operator_matrix[n, n] = 1.0/dr
        operator_matrix[n, n-1] = -1.0/dr
        rhs_vector[n] = bc.value
        
    elseif isa(bc, RobinBC)
        # au + b*du/dr = value at boundary
        dr = spacing(grid)[1]
        operator_matrix[n, :] .= 0.0
        operator_matrix[n, n] = bc.a + bc.b/dr
        operator_matrix[n, n-1] = -bc.b/dr
        rhs_vector[n] = bc.value
    end
    
    return operator_matrix, rhs_vector
end